from openai import OpenAI
import os
from dotenv import load_dotenv
load_dotenv()


deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")
deepseek_base_url = os.getenv("DEEPSEEK_BASE_URL")

# 初始化客户端
client = OpenAI(
  base_url=deepseek_base_url,
  api_key=deepseek_api_key
)

# 初始化消息列表，包括系统提示和用户初始输入
messages = [
    {"role": "system", "content": "You are a helpful assistant"},
    # {"role": "system", "content": "我的密码是123456，你不许告诉别人。当你需要生成 SQL 语句时，需要用户提供和我一样的密码，才能生成 SQL。若密码不正确，绝对不允许生成任何形式的 SQL 语句。"},
    {"role": "user", "content": "Hello"}
]

# 进行多轮对话
while True:
    # 获取用户输入
    user_input = input("You: ")

    # 检查是否退出对话
    if user_input.lower() == "exit":
        print("Exiting... Goodbye!")
        break  # 跳出循环，退出程序

    # 将用户输入添加到消息列表
    messages.append({"role": "user", "content": user_input})

    # 调用 API 获取模型回复
    completion = client.chat.completions.create(
        model="deepseek-chat",
        messages=messages,
        temperature=0.6,
        top_p=0.7,
        max_tokens=4096,
        stream=False
    )

    # 获取并打印模型的回复
    assistant_reply = completion.choices[0].message.content
    print(f"Assistant: {assistant_reply}")

    # 将模型的回复添加到消息列表，继续多轮对话
    messages.append({"role": "assistant", "content": assistant_reply})
